import numpy as np
import scipy.special as sc
import scipy.optimize as op
from chaco.api import *
from traits.api import *
from traitsui.api import *
from chaco.tools.api import LineInspector, PanTool, ZoomTool
from enable.api import ComponentEditor
import time
from PyQt5.QtCore import QThread
import meshio


def importMesh(file):
    mesh = meshio.read(file)
    v = np.array(mesh.points)
    f = np.zeros((0, 3), dtype=int)
    for i in mesh.cells:
        if i.type == 'triangle':
            f = np.vstack((f, np.array(i.data)))
    return v, f


def extractEdges(self, f):
    nFacet = f.shape[0]
    f = np.vstack((f, np.roll(f, -1, axis=-1), np.roll(f, -2, axis=-1)))
    e = np.equal.outer(f[:, 1], f[:, 2])
    e = e & e.T & np.triu(np.ones((f.shape[0], f.shape[0]), dtype=bool))
    e = e.nonzero()
    edges = np.hstack((f[e[0], :], f[e[1], 0:1]))
    edgesFacet = e.T % nFacet
    return edges, edgesFacet


class Form(HasTraits):
    frequencies = List(Float, [3, 10])
    elevations = List(Float, [0, 10])
    azimuthStart = Float(0)
    azimuthStop = Float(45)
    azimuthNumber = Int(91)
    zabo = Float(91)
    sunhao = Float(82)
    # po = Enum('HH', 'VV')
    plot = Instance(Plot)
    progress = Str('0%')
    buttonParameter = Button('参数设置')
    buttonCalculate = Button('计算')
    buttonStop = Button('停止')
    buttonExtractEdges = Button('提取边缘')
    method = Enum('PO', 'PTD', 'PO+PTD')
    thread = QThread()
    atmosphere=Float(5.6)
    atmosphere1=Float(5.63)

    def _buttonExtractEdges_fired(self):
        self.importMesh()
        f = self.f
        f = np.vstack((f, np.roll(f, -1, axis=-1), np.roll(f, -2, axis=-1)))
        e = np.equal.outer(f[:, 1], f[:, 2])
        e = e & e.T & np.triu(np.ones((f.shape[0], f.shape[0]), dtype=bool))
        e = e.nonzero()
        self.edges = np.hstack((f[e[0], :], f[e[1], 0:1]))
        self.edgesFacet = e.T

    def extractEdges(self, f):
        nFacet = f.shape[0]
        f = np.vstack((f, np.roll(f, -1, axis=-1), np.roll(f, -2, axis=-1)))
        e = np.equal.outer(f[:, 1], f[:, 2])
        e = e & e.T & np.triu(np.ones((f.shape[0], f.shape[0]), dtype=bool))
        e = e.nonzero()
        edges = np.hstack((f[e[0], :], f[e[1], 0:1]))
        edgesFacet = e.T % nFacet
        return edges, edgesFacet

    def importMesh(self, file):
        mesh = meshio.read(file)
        v = np.array(mesh.points)
        f = np.zeros((0, 3), dtype=int)
        for i in mesh.cells:
            if i.type == 'triangle':
                f = np.vstack((f, np.array(i.data)))
        return v, f

        f = np.vstack((f, np.roll(f, -1, axis=-1), np.roll(f, -2, axis=-1)))
        e = np.equal.outer(f[:, 1], f[:, 2])
        e = e & e.T & np.triu(np.ones((f.shape[0], f.shape[0]), dtype=bool))
        e = e.nonzero()
        self.e = np.hstack((f[e[0], :], f[e[1], 0:1]))
        e = e
        self.ef = e.T

    def calculate(self):
        v, f = importMesh('delete.stl')
        edges, edgesFacet = extractEdges(f)

        frequencies = np.array(self.frequencies)
        elevations = np.array(self.elevations)
        azimuths = np.linspace(self.azimuthStart, self.azimuthStop, self.azimuthNumber)
        self.rcs = np.full((2, frequencies.size, elevations.size, azimuths.size), np.nan)
        np.seterr(divide='ignore')
        # v = self.root.model.vertices
        # f = self.root.model.facets
        v = np.concatenate((v[f[:, 0]][:, np.newaxis, :], v[f[:, 1]][:, np.newaxis, :], v[f[:, 2]][:, np.newaxis, :]), axis=1)
        e = np.roll(v, -1, axis=1) - v
        ec = (np.roll(v, -1, axis=1) + v) / 2
        n = np.cross(e[:, 0], e[:, 1])
        a = np.linalg.norm(n, axis=-1) / 2
        n = n / a[:, np.newaxis] / 2
        eps = 1e-8
        totalNumber = frequencies.size * elevations.size * azimuths.size
        for ff, frequency in enumerate(frequencies):
            k = 2 * np.pi / .3 * frequency
            for ee, elevation in enumerate(elevations):
                cosel = np.cos(elevation * np.pi / 180)
                sinel = np.sin(elevation * np.pi / 180)
                for aa, azimuth in enumerate(azimuths):
                    cosaz = np.cos(azimuth * np.pi / 180)
                    sinaz = np.sin(azimuth * np.pi / 180)
                    i = -np.array([cosel * cosaz, cosel * sinaz, sinel])
                    ni = n @ i
                    perpendicular = ni <= -1 + eps
                    oblique = (ni < -eps) & ~perpendicular
                    w = -2 * i
                    poPerpendicular = np.einsum('i,i,i->', ni[perpendicular], np.exp(1j * k * v[perpendicular, 0, :] @ w), a[perpendicular]) * 1j * k / np.sqrt(np.pi)
                    ncw = np.cross(n[oblique], w, axis=-1)
                    ncw2 = (ncw * ncw).sum(-1)
                    poOblique = np.einsum('i,ik,ijk,ij->', ni[oblique] / ncw2, ncw, e[oblique], np.exp(1j * k * ec[oblique] @ w) * np.sinc(k * e[oblique] @ w / 2 / np.pi)) / np.sqrt(np.pi)
                    po = poPerpendicular + poOblique
                    self.rcs[0, ff, ee, aa] = 20 * np.log10(np.abs(po))
                    self.rcs[1, ff, ee, aa] = 20 * np.log10(np.abs(po))
                    self.plot.data.update_data(
                        x=azimuths[:aa + 1],
                        hh=self.rcs[0, ff, ee, :aa + 1],
                        vv=self.rcs[1, ff, ee, :aa + 1],
                    )
                    self.plot.title = f'Frequency={frequency}GHz,Elevation={elevation} degree'
                    self.progress = f'{((ff * elevations.size + ee) * azimuths.size + aa + 1) / totalNumber:<.1%}'
                    time.sleep(.1)

    def _buttonParameter_fired(self):
        self.configure_traits(view='parameter_view')

    def _buttonCalculate_fired(self):
        # self.calculate()
        self.thread.start()

    def _buttonStop_fired(self):
        self.thread.terminate()
        self.thread.wait()

    def traits_init(self):
        self.thread.run = self.calculate
        self.plot = Plot(ArrayPlotData())
        hh = np.arange(10, dtype=float)
        hh[5] = np.nan
        self.plot.data.update_data(
            x=np.arange(10),
            hh=hh,
            vv=np.arange(10) + 1,
        )
        self.plot.title = 'PO'
        self.plot.x_axis.title = "Azimuth/deg"
        self.plot.y_axis.title = "RCS/dBsm"
        self.plot.plot(("x", "hh"), name='HH', color='auto')
        self.plot.plot(("x", "vv"), name='VV', color='auto')
        legend = Legend()
        legend.plots = self.plot.plots
        self.plot.overlays.append(legend)
        # self.plot.tools.append(
        #     PanTool(self.plot, constrain_key="shift")
        # )
        # self.plot.overlays.append(
        #     ZoomTool(component=self.plot, tool_mode="box", always_on=False)
        # )

    traits_view = View(
        UItem(
            "plot",
            editor=ComponentEditor(),
            resizable=True,
        ),
        HGroup(
            UItem('buttonParameter'),
            UItem('buttonCalculate'),
            UItem('buttonStop'),
            UItem('buttonExtractEdges'),
            Item('zabo', label='杂波'),
            Item('sunhao', label='损耗'),
            # Item('po', label='极化', style='custom'),
            UItem('progress'),
            Item('atmosphere',label='大气损耗'),
            Item('atmosphere1',label='大气损耗1'),
            spring
        ),
        width=1200,
        height=900,
        resizable=True,
    )
    parameter_view = View(
        Group(
            UItem('frequencies'),
            label='频率/GHz',
            show_border=True,
        ),
        Group(
            UItem('elevations'),
            label='高低角/°',
            show_border=True,
        ),
        Item('azimuthStart', label='方位角起始/°'),
        Item('azimuthStop', label='方位角终止/°'),
        Item('azimuthNumber', label='方位角数量'),
        height=600,
        resizable=True,
        buttons=OKCancelButtons,
    )


form = Form()
if __name__ == '__main__':
    form.configure_traits()
